from pylab import *
from camori_aux import harris_vdirs, directional_error
from quaternion import *
import scipy
import scipy.spatial

# def edgel_distance(e1, e2):
#     return (((e2[:2]-e1[:2])**2).sum()**0.5 /
#             np.abs(e2[2]*e1[3] + e2[3]*e1[2]))
# def edgel_distance(e1, e2):
#     v = e2[:2]-e1[:2]
#     vn = (v**2).sum()**0.5
#     v = v / vn

#     k = 1.0
#     return vn * k / (k + np.abs(e1[2]*v[1]+e1[3]*v[0])
#                      + np.abs(e2[2]*v[1]+e2[3]*v[0]))

def numpy_directional_error(e1, e2):
    v = e2[:2]-e1[:2]
    vn = (v**2).sum()**0.5
    v = v / vn
    sig_a = 0.5

    # pa1 = exp(-((e1[2]*v[0]+e1[3]*v[1])/sig_a)**2)
    # pa2 = exp(-((e2[2]*v[0]+e2[3]*v[1])/sig_a)**2)
    # return pa1*pa2
    pa1 = (e1[2]*v[0]+e1[3]*v[1])**2
    pa2 = (e2[2]*v[0]+e2[3]*v[1])**2
    return pa1+pa2


if __name__ == '__main__':
    Ned = 10000
    edgels = zeros((Ned,4), dtype=np.float32)
    edgels[:,:2] = rand(Ned, 2) * 256

    q = random_quaternion()
    #q = Quat(1.0,0,0,0)

    hh = harris_vdirs(128,128,128,0.0,q.q,edgels)
    lab = random_integers(0,2,Ned)
    for n in range(Ned):
        if n%4 == 0:
            edgels[n,2] = hh[0,n,1]
            edgels[n,3] = -hh[0,n,0]
        if n%4 == 1:
            edgels[n,2] = hh[1,n,1]
            edgels[n,3] = -hh[1,n,0]
        if n%4 == 2:
            edgels[n,2] = hh[2,n,1]
            edgels[n,3] = -hh[2,n,0]
        if n%4 == 3:
            edgels[n,2:] = Quat(sin(rand()*2*pi),0,0).q[:2]

    # edgels[0,0] = 128
    # edgels[0,1] = 128
    # edgels[0,2] = 0.0
    # edgels[0,3] = 1.0

    catch_all = False

    kdt = scipy.spatial.cKDTree( edgels[:,:2] )
    templates = {}    
    Nned = 0
    gfil_d = 10.0
    esel = np.zeros(Ned, dtype=int)
    count = np.zeros(Ned, dtype=int)
    ## If we "gotta catch them all", we just count for each edgel how
    ## many neighbours it has. But the neighborhood doesnp t take just
    ## in account the dstance, that is what the kdt uses to give us
    ## the preliminary neighbours list. We also look at the angular
    ## error to further discard some of these initial neighboring
    ## edgels.
    if catch_all:
        for n in range(0,Ned):
            px,py,ux,uy = edgels[n]
            dist, idx = kdt.query(np.array([[px,py]]), 100, eps=0, p=2, distance_upper_bound=gfil_d)
            hits = 0
            for d, m in zip(dist[0],idx[0]):
                ## Discard the last position of the output array from the query.
                if m == Ned:
                    continue
                ## Calculate the angular error 
                err_a, err_d = directional_error(edgels, m, n)
                if err_a < 0.9:
                    hits += 1
            if catch_all:
                templates[n] = hits
        ## templates now contains hoe many distance+angule neighbors each
        ## edgel has, and we can use this to discard "lonely", edgels with
        ## few similar ones around.

        ## This alternative is in case we want to perform a
        ## homogeneization of the density of observations. We discard
        ## edgels that already have a similar one in the output list.
    else:
        for n in range(0,Ned):
            px,py,ux,uy = edgels[n]
            dist, idx = kdt.query(np.array([[px,py]]), 100, eps=0, p=2, distance_upper_bound=gfil_d)
            hits = 0
            for d, m in zip(dist[0],idx[0]):
                ## Discard the last position of the output array from the query.
                if m == Ned:
                    continue
                ## We only count hits in the edgels already in the templates list.
                if not m in templates:
                    continue
                err_a, err_d = directional_error(edgels, m, n)
                if err_a < 0.9:
                    ## strengthen existing neighbor templates
                    templates[m] += 1
                    hits += 1
            ## No hits, this can become a new template. We might use
            ## "hits" here, but I am not sre exactly how right
            ## now. Better to start with 0 and let new edgels vote for
            ## it.
            if hits == 0:
                templates[n] = 0
        ## templates now contain a list of edgels that "dominate" it
        ## own regions, no two templates are close to each other, and
        ## their counts gve an estimate of how many of the original
        ## observations are similar to it. Now we can discard the
        ## weaker observations, like previously, bu we also
        ## avoidhaving too many instances of very similar
        ## observations.

    ## Now we pick up the final likst of output edgels, possibly
    ## dicarding "weak" observations.
    esel = [k for k in templates if templates[k]>=2]
    Nned = len(esel)

    ion()
    figure(1, figsize=(12,6))
    vv = 4*array([-1,1])

    # subplot(1,2,1)
    # for n in range(Ned):
    #     px,py,ux,uy=edgels[n]
    #     plot(px+uy*vv, py-ux*vv, 'r-', alpha=0.7)
    # axis('equal')
    # axis([0,256,256,0])
    # subplot(1,2,2)
    # for m in range(Nned):
    #     n = esel[m]
    #     px,py,ux,uy=edgels[n]
    #     plot(px+uy*vv, py-ux*vv, 'r-', alpha=0.7)
    # axis('equal')
    # axis([0,256,256,0])

    for n in range(Ned):
        px,py,ux,uy=edgels[n]
        if n%4 == 0:
            subplot(2,4,1)
            plot(px+uy*vv, py-ux*vv, 'r-', alpha=0.7)
        if n%4 == 1:
            subplot(2,4,2)
            plot(px+uy*vv, py-ux*vv, 'g-', alpha=0.7)
        if n%4 == 2:
            subplot(2,4,3)
            plot(px+uy*vv, py-ux*vv, 'b-', alpha=0.7)
        if n%4 == 3:
            subplot(2,4,4)
            plot(px+uy*vv, py-ux*vv, 'k-', alpha=0.2)

    for m in range(Nned):
        n = esel[m]
        px,py,ux,uy=edgels[n]
        if n%4 == 0:
            subplot(2,4,5)
            plot(px+uy*vv, py-ux*vv, 'r-', alpha=0.7)
        if n%4 == 1:
            subplot(2,4,6)
            plot(px+uy*vv, py-ux*vv, 'g-', alpha=0.7)
        if n%4 == 2:
            subplot(2,4,7)
            plot(px+uy*vv, py-ux*vv, 'b-', alpha=0.7)
        if n%4 == 3:
            subplot(2,4,8)
            plot(px+uy*vv, py-ux*vv, 'k-', alpha=0.2)

    for k in range(8):
        subplot(2,4,k+1)
        axis('equal')
        axis([-10.0,266,266,-10])

        
    print [len([k for k in templates if (templates[k]%4)==z]) for z in range(4)]
